package demoflink.state;

import demoflink.entity.WaterSensor;
import demoflink.function.WaterSensorMapFunction;
import org.apache.flink.api.common.state.BroadcastState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.ReadOnlyBroadcastState;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.contrib.streaming.state.EmbeddedRocksDBStateBackend;
import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
import org.apache.flink.streaming.api.datastream.BroadcastConnectedStream;
import org.apache.flink.streaming.api.datastream.BroadcastStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction;
import org.apache.flink.util.Collector;

public class OperatorBroadcastStateDemo {
    public static void main(String[] args) {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(2);

        //状态后端的存储代码
        HashMapStateBackend hashMapStateBackend = new HashMapStateBackend();
        env.setStateBackend(hashMapStateBackend);
//        EmbeddedRocksDBStateBackend embeddedRocksDBStateBackend = new EmbeddedRocksDBStateBackend();
//        env.setStateBackend(embeddedRocksDBStateBackend);



        SingleOutputStreamOperator<WaterSensor> sensorDs = env.socketTextStream("node1", 7777)
                .map(new WaterSensorMapFunction());
        DataStreamSource<String> configDs = env.socketTextStream("node1", 8888);
        MapStateDescriptor<String, Integer> descriptor = new MapStateDescriptor<>("broadcast-state", Types.STRING, Types.INT);
        BroadcastStream<String> broadcast = configDs.broadcast(descriptor);
        BroadcastConnectedStream<WaterSensor, String> sensorBCS = sensorDs.connect(broadcast);
        sensorBCS.process(new BroadcastProcessFunction<WaterSensor, String, String>() {
            @Override
            public void processElement(WaterSensor waterSensor, BroadcastProcessFunction<WaterSensor, String, String>.ReadOnlyContext context, Collector<String> collector) throws Exception {
                ReadOnlyBroadcastState<String, Integer> broadcastState = context.getBroadcastState(descriptor);
                Integer threshold = broadcastState.get("threshold");
                if (threshold==null) {
                    threshold = 0;
                }
                if (waterSensor.getVc()>threshold) {
                    collector.collect(waterSensor+",水位超过指定的阈值："+threshold+"!");
                }
            }

            /**
             * 广播的配置流的处理方法
             * @param s
             * @param context
             * @param collector
             * @throws Exception
             */
            @Override
            public void processBroadcastElement(String s, BroadcastProcessFunction<WaterSensor, String, String>.Context context, Collector<String> collector) throws Exception {
                BroadcastState<String, Integer> broadcastState = context.getBroadcastState(descriptor);
                broadcastState.put("threshold",Integer.parseInt(s));
            }
        }).print();

    }
}
